# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: inkyu

from coin.base.book import BookBuilderBase
from coin.exchange.huobi.kr_rest.product import HuobiProduct


class BookEntry(object):
  def __init__(self, price, qty):
    self.price = float(price)
    self.qty = float(qty)


class BookBuilder(BookBuilderBase):
  ProductType = HuobiProduct

  def __init__(self,
               native_symbol,
               current_datetime=None,
               book_builder_observer=None,
               crop_book_snapshot_depth=None):
    super().__init__("Huobi",
                     native_symbol,
                     current_datetime,
                     book_builder_observer=book_builder_observer)
    self._crop_book_snapshot_depth = crop_book_snapshot_depth
    self._last_seq_num = -1
    self._seq_num_successive = False

  def snapshot(self, bids: list, asks: list):
    self.clear_book()

    if self._crop_book_snapshot_depth is not None:
      assert self._crop_book_snapshot_depth >= 10, self._crop_book_snapshot_depth
      bids = sorted(bids, key=lambda x: x[0], reverse=True)[0:self._crop_book_snapshot_depth]
      asks = sorted(asks, key=lambda x: x[0], reverse=False)[0:self._crop_book_snapshot_depth]

    for data in bids:
      price, qty = [float(value) for value in data]
      self._insert_bid(price, qty, insert_only=True)
    for data in asks:
      price, qty = [float(value) for value in data]
      self._insert_ask(price, qty, insert_only=True)

  def _insert_bid(self, price, qty, insert_only=False):
    assert price > 0
    assert qty > 0
    entry = BookEntry(price, qty)
    if insert_only:
      assert price not in self._bid_price_map
      self._bid_price_map[price] = entry
    else:
      self._bid_price_map[price] = entry
      if self.ask0() is None:
        return
      bid_price = self.bid0().price
      while self.ask0() is not None and bid_price >= self.ask0().price:
        ask_price = self.ask0().price
        self._delete_ask(ask_price)

  def _insert_ask(self, price, qty, insert_only=False):
    assert price > 0
    assert qty > 0
    entry = BookEntry(price, qty)
    if insert_only:
      assert price not in self._ask_price_map
      self._ask_price_map[price] = entry
    else:
      self._ask_price_map[price] = entry
      if self.bid0() is None:
        return
      ask_price = self.ask0().price
      while self.bid0() is not None and ask_price <= self.bid0().price:
        bid_price = self.bid0().price
        self._delete_bid(bid_price)

  def _delete_bid(self, price):
    if price in self._bid_price_map:
      del self._bid_price_map[price]

  def _delete_ask(self, price):
    if price in self._ask_price_map:
      del self._ask_price_map[price]

  def update_diff(self, bids: list, asks: list, seq_num: int, prev_seq_num: int):
    if self._last_seq_num == -1:
      self._last_seq_num = seq_num
      self._seq_num_successive = True
    elif self._seq_num_successive:
      if prev_seq_num == self._last_seq_num:
        self._last_seq_num = seq_num
      else:
        self._seq_num_successive = False

    for data in bids:
      price = float(data[0])
      qty = float(data[1])
      if qty > 0:
        self._insert_bid(price, qty, insert_only=False)
      else:
        self._delete_bid(price)

    for data in asks:
      price = float(data[0])
      qty = float(data[1])
      if qty > 0:
        self._insert_ask(price, qty, insert_only=False)
      else:
        self._delete_ask(price)
